from langchain import hub
import akshare as ak
from langchain_openai import ChatOpenAI
from langchain_deepseek import ChatDeepSeek
from langgraph.prebuilt import create_react_agent
from langchain.tools import tool
from langchain.output_parsers import StructuredOutputParser
from langchain.output_parsers import ResponseSchema
import pypandoc
import os

@tool
def get_stock_income(stock_code: str) -> str:
    """ 
    获取股票利润表数据 
    参数:
    stock_code (str): 股票代码
 
    返回:
    str: 股票利润表数据。
    """
    income_df = ak.stock_financial_report_sina(stock=stock_code, symbol="利润表")
    return str(income_df.head(20).to_dict(orient="records"))
    
@tool
def get_stock_balance(stock_code: str) -> str:
    """
    获取股票资产负债表数据
    参数:
    stock_code (str): 股票代码
 
    返回:
    str: 股票资产负债表数据。
    """
    balance_df = ak.stock_financial_report_sina(stock=stock_code, symbol="资产负债表")   
    r = str(balance_df.head(20).to_dict(orient="records"))
    return r

@tool
def add(a:int, b:int) -> int:
    """Add two numbers together"""
    return a + b

@tool
def multiply(a:int, b:int) -> int:
    """Multiply two numbers together"""
    return a * b

def analyze_stock_tools(sec_code:str) -> dict:
    # 从环境变量中获取API密钥
    try:
        api_key = os.environ['API_KEY']
        api_base = os.environ['API_BASE']
        model_name = os.environ['MODEL_NAME']
    except KeyError:
        raise ValueError('DEEPSEEK_API_KEY未在环境变量中配置')

    llm = ChatDeepSeek(
        model=model_name,
        api_key=api_key,
        api_base=api_base,
        timeout=1800,
    )
    # Choose the LLM that will drive the agent
    tools = [get_stock_income, get_stock_balance]


    zhpj_schema = ResponseSchema(name="综合评价", description="给出该股票的综合评价，1-5星，只要给出数字")
    cwjk_schema = ResponseSchema(name="财务状况是否健康", description="给出该股票的财务健康状况，只要给出是或否")
    tzjz_schema = ResponseSchema(name="是否具有投资价值", description="给出该股票的投资价值，只要给出是或否")
    xxpj_schema = ResponseSchema(name="详细评价", description="给出该股票的基本面的详细评价")
    response_schemas = [zhpj_schema, cwjk_schema, tzjz_schema, xxpj_schema]

    output_parser = StructuredOutputParser.from_response_schemas(response_schemas)

    prompt = """
    你是一个专业的股票分析师，请根据用户给出的股票代码使用tools工具来获取利润表、资产负债表数据，然后对其进行分析。
    请按以下结构输出：
    综合评价:给出该股票的综合评价，1-5星，步进0.5，只要给出数字
    财务状况是否健康:给出该股票的财务健康状况，只要给出是或否
    是否具有投资价值:给出该股票的投资价值，只要给出是或否
    详细评价:给出该股票的基本面的详细评价,markdown格式,多方面阐述评价的结论能图表结合更佳。

    ```json
    {
        "综合评价": string   
        "财务状况是否健康": string   
        "是否具有投资价值": string   
        "详细评价": string  
    }
    ```
    """
    agent_executor = create_react_agent(llm, tools, prompt=prompt)
    result =  agent_executor.invoke({"messages": [("user", f"{sec_code}")]})
    print("="*100)
    print(type(result))
    print("-"*100)
    print(result)
    print("="*100)
    print(result["messages"][-1].content)
    rjson = output_parser.parse(result["messages"][-1].content)
    filename = 'd:\\pdf\\'+sec_code+'.pdf'
    output = pypandoc.convert_text(rjson["详细评价"], 'pdf', 'markdown', outputfile=filename)
    print("="*100)
    print(output)
    return rjson
